import numpy as np
from typing import NamedTuple, Sequence

from .hand import NUM_JOINTS_PER_HAND


class SingleHandPose(NamedTuple):
    """
    A hand pose is composed of two fields:
    1) joint angles where # joints == # DoFs
    2) root-to-world rigid wrist transformation
    """

    joint_angles: np.ndarray = np.zeros(NUM_JOINTS_PER_HAND, dtype=np.float32)
    wrist_xform: np.ndarray = np.eye(4, dtype=np.float32)
    # hand_confidence: float = 1.0
    hand_confidence: np.ndarray = np.array([0.0], dtype=np.float32)


class BatchHandPose(NamedTuple):
    """
    A batched hand pose comprises:
    - joint_angles: array of shape (B, NUM_JOINTS_PER_HAND)
    - wrist_xform: array of shape (B, 4, 4)
    - hand_confidence: array of shape (B,)
    """

    joint_angles: np.ndarray
    wrist_xform: np.ndarray
    hand_confidence: np.ndarray


def batch_single_hand(pose: SingleHandPose, batch_size: int) -> BatchHandPose:
    """
    Replicate a single hand pose across a batch dimension.
    :param pose: SingleHandPose to replicate
    :param batch_size: number of duplicates
    """
    # Expand dims and tile
    ja = np.broadcast_to(pose.joint_angles[None, ...], (batch_size, *pose.joint_angles.shape)).copy()
    wx = np.broadcast_to(pose.wrist_xform[None, ...], (batch_size, 4, 4)).copy()
    hc = np.full((batch_size,), pose.hand_confidence, dtype=np.float32)
    return BatchHandPose(joint_angles=ja, wrist_xform=wx, hand_confidence=hc)


def stack_hand_poses(poses: Sequence[SingleHandPose]) -> BatchHandPose:
    """
    Stack a sequence of SingleHandPose into a batched representation.
    :param poses: iterable of SingleHandPose, length B
    """
    ja_list = []
    wx_list = []
    hc_list = []
    for p in poses:
        ja_list.append(p.joint_angles)
        wx_list.append(p.wrist_xform)
        hc_list.append(p.hand_confidence)

    ja = np.stack(ja_list, axis=0)
    wx = np.stack(wx_list, axis=0)
    hc = np.array(hc_list, dtype=np.float32)
    return BatchHandPose(joint_angles=ja, wrist_xform=wx, hand_confidence=hc)
